package top.rainbowecho.gateway.security.jwt;

import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.core.userdetails.User;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import top.rainbowecho.common.util.JwtTokenKey;
import top.rainbowecho.gateway.mapper.UsersMapper;
import top.rainbowecho.model.field.UsersField;
import top.rainbowecho.model.percona.Users;

import java.util.HashMap;
import java.util.Map;

/**
 * 用于对token内容进行自定义增强
 *
 * @author rainbow
 * @since 2020/10/10 15:19
 */
@Slf4j
public class JWTokenEnhancer implements TokenEnhancer {
    private UsersMapper usersMapper;

    public void setUsersMapper(UsersMapper usersMapper) {
        this.usersMapper = usersMapper;
    }

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken oAuth2AccessToken, OAuth2Authentication oAuth2Authentication) {
        Map<String, Object> info = new HashMap<>();
        // 拿到认证用户，查询数据库
        User authenticationUser = (User) oAuth2Authentication.getUserAuthentication().getPrincipal();
        QueryWrapper<Users> userQueryWrapper = new QueryWrapper<Users>().eq(UsersField.USER_NAME, authenticationUser.getUsername());
        Users user = this.usersMapper.selectOne(userQueryWrapper);

        // 将用户信息放到token中
        info.put(JwtTokenKey.USER_ID, user.getId());
        ((DefaultOAuth2AccessToken) oAuth2AccessToken).setAdditionalInformation(info);
        return oAuth2AccessToken;
    }
}
